# -*- coding: utf-8 -*-
# @Time    : 2021/10/14 15:59
# @Author  : huangwei
# @File    : get_cell.py
# @Software: PyCharm
import cv2
import numpy as np
from skimage import measure

from table_method import get_line, draw_lines, sqrt, line_line, roll_min_rect_box
from table_net import table_net
import tensorflow as tf

config = tf.compat.v1.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.5
config.gpu_options.allow_growth = True
session = tf.compat.v1.InteractiveSession(config=config)

tableModeLinePath = 'models/table-line.h5'
model = table_net((None, None, 3), 2)
model.load_weights(tableModeLinePath)


def resize( img, size ):
    img_h, img_w = img.shape[:2]
    w, h = size

    ratio = min(w / img_w, h / img_h)
    new_w = int(img_w * ratio)
    new_h = int(img_h * ratio)

    resize_img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_CUBIC)
    cv2.imwrite('temp_path/resize_img.png', resize_img)

    fill_value = [128, 128, 128]
    boxed_image = np.zeros((h, w, 3), dtype=np.uint8)
    boxed_image[:] = fill_value
    boxed_image[:new_h, :new_w, :] = resize_img

    # 分别为长、宽变化的比例
    fx = new_w / img_w
    fy = new_h / img_h

    return boxed_image, fx, fy


def adjust_lines( RowsLines, ColsLines, alph=50 ):
    ##调整line，合并相交线
    nrow = len(RowsLines)
    ncol = len(ColsLines)
    newRowsLines = []
    newColsLines = []
    for i in range(nrow):

        x1, y1, x2, y2 = RowsLines[i]
        cx1, cy1 = (x1 + x2) / 2, (y1 + y2) / 2
        for j in range(nrow):
            if i != j:
                x3, y3, x4, y4 = RowsLines[j]
                cx2, cy2 = (x3 + x4) / 2, (y3 + y4) / 2
                if (x3 < cx1 < x4 or y3 < cy1 < y4) or (x1 < cx2 < x2 or y1 < cy2 < y2):
                    continue
                else:
                    r = sqrt((x1, y1), (x3, y3))
                    if r < alph:
                        newRowsLines.append([x1, y1, x3, y3])
                    r = sqrt((x1, y1), (x4, y4))
                    if r < alph:
                        newRowsLines.append([x1, y1, x4, y4])

                    r = sqrt((x2, y2), (x3, y3))
                    if r < alph:
                        newRowsLines.append([x2, y2, x3, y3])
                    r = sqrt((x2, y2), (x4, y4))
                    if r < alph:
                        newRowsLines.append([x2, y2, x4, y4])

    for i in range(ncol):
        x1, y1, x2, y2 = ColsLines[i]
        cx1, cy1 = (x1 + x2) / 2, (y1 + y2) / 2
        for j in range(ncol):
            if i != j:
                x3, y3, x4, y4 = ColsLines[j]
                cx2, cy2 = (x3 + x4) / 2, (y3 + y4) / 2
                if (x3 < cx1 < x4 or y3 < cy1 < y4) or (x1 < cx2 < x2 or y1 < cy2 < y2):
                    continue
                else:
                    r = sqrt((x1, y1), (x3, y3))
                    if r < alph:
                        newColsLines.append([x1, y1, x3, y3])
                    r = sqrt((x1, y1), (x4, y4))
                    if r < alph:
                        newColsLines.append([x1, y1, x4, y4])

                    r = sqrt((x2, y2), (x3, y3))
                    if r < alph:
                        newColsLines.append([x2, y2, x3, y3])
                    r = sqrt((x2, y2), (x4, y4))
                    if r < alph:
                        newColsLines.append([x2, y2, x4, y4])

    return newRowsLines, newColsLines


def get_table_line( pred_line, axis, alph ):
    # axis=0 横线
    # axis=1 竖线
    lables = measure.label(pred_line > 0, connectivity=2)
    regions = measure.regionprops(lables)

    if axis == 1:
        lines = [get_line(line.coords) for line in regions if line.bbox[2] - line.bbox[0] > alph]
    else:
        lines = [get_line(line.coords) for line in regions if line.bbox[3] - line.bbox[1] > alph]

    return lines


def get_lines( img, input_size, prob=0.5, alph1=10, alph2=15, row=30, col=30 ):
    resize_img, fx, fy = resize(img, input_size)

    pred = model.predict(np.array([np.array(resize_img) / 255.0]))
    pred = pred[0]

    rows = pred[..., 0] > prob
    cols = pred[..., 1] > prob
    rows = rows.astype(int)
    cols = cols.astype(int)

    row_lines = get_table_line(rows, axis=0, alph=row)
    col_lines = get_table_line(cols, axis=1, alph=col)

    tmp = np.zeros(input_size, dtype='uint8')
    line_img = draw_lines(tmp, row_lines + col_lines, color=255, line_width=2)
    cv2.imwrite('temp_path/bb.png', line_img)

    # 最大的外接线
    ccolbox = []
    crowlbox = []
    if len(row_lines) > 0:
        row_lines = np.array(row_lines)
        row_lines[:, [0, 2]] = row_lines[:, [0, 2]] / fx
        row_lines[:, [1, 3]] = row_lines[:, [1, 3]] / fy
        xmin = row_lines[:, [0, 2]].min()
        xmax = row_lines[:, [0, 2]].max()
        ymin = row_lines[:, [1, 3]].min()
        ymax = row_lines[:, [1, 3]].max()
        # 最大的外接线
        ccolbox = [[xmin, ymin, xmin, ymax], [xmax, ymin, xmax, ymax]]
        row_lines = row_lines.tolist()

    if len(col_lines) > 0:
        col_lines = np.array(col_lines)
        col_lines[:, [0, 2]] = col_lines[:, [0, 2]] / fx
        col_lines[:, [1, 3]] = col_lines[:, [1, 3]] / fy

        xmin = col_lines[:, [0, 2]].min()
        xmax = col_lines[:, [0, 2]].max()
        ymin = col_lines[:, [1, 3]].min()
        ymax = col_lines[:, [1, 3]].max()
        col_lines = col_lines.tolist()
        crowlbox = [[xmin, ymin, xmax, ymin], [xmin, ymax, xmax, ymax]]

    # row_lines += crowlbox
    # col_lines += ccolbox

    # 同一种类型的线是否重合进行合并
    row_box, col_box = adjust_lines(row_lines, col_lines, alph=alph1)

    row_lines += row_box
    col_lines += col_box

    # 横线数线是否相交进行连接，是的话则延长至交点
    for i in range(len(row_lines)):
        for j in range(len(col_lines)):
            row_lines[i] = line_line(row_lines[i], col_lines[j], alph2)
            col_lines[j] = line_line(col_lines[j], row_lines[i], alph2)

    # 去除不平行的线
    # 去除长度不够的线

    return row_lines, col_lines


file_path = "images/simple3.png"
img = cv2.imread(file_path)
input_size = (1024, 1024)
row_lines, col_lines = get_lines(img, input_size)

tmp = np.zeros(img.shape[:2], dtype='uint8')
line_img = draw_lines(tmp, row_lines + col_lines, color=255, line_width=2)
cv2.imwrite('temp_path/aa.png', line_img)

labels = measure.label(line_img < 255, connectivity=2)

labels1 = np.copy(labels)
for m in range(labels1.shape[0]):
    for n in range(labels1.shape[1]):
        if labels1[m][n] != 0:
            labels1[m][n] = 200
cv2.imwrite('temp_path/ee.png', labels1)

regions = measure.regionprops(labels)
print("regions:",len(regions))

# 取出符合要求的外接矩形
rect_boxes = roll_min_rect_box(regions, line_img.shape)
boxes = np.array(rect_boxes)
# 将检测出的框画在图片上
from method import draw_box

for box in boxes:
    draw_box(img, box)
cv2.imwrite("temp_path/box.png", img)
print("检测出的表格框数量为：", len(boxes))
